import os
import pymysql
from jinja2 import Environment, FileSystemLoader


class CodeGenerator:
    def __init__(self, db_config):
        self.db_config = db_config
        self.env = Environment(loader=FileSystemLoader('templates'))
        self.connection = None
        # MySQL数据类型到Java类型的映射
        self.type_mapping = {
            'int': 'Integer',
            'bigint': 'Long',
            'varchar': 'String',
            'char': 'String',
            'text': 'String',
            'longtext': 'String',
            'datetime': 'Date',
            'date': 'Date',
            'timestamp': 'Date',
            'decimal': 'BigDecimal',
            'double': 'Double',
            'float': 'Float',
            'tinyint': 'Integer',
            'smallint': 'Integer',
            'mediumint': 'Integer',
            'bit': 'Boolean',
            'json': 'String',
            'enum': 'String',
            'set': 'String',
            'binary': 'byte[]',
            'varbinary': 'byte[]',
            'blob': 'byte[]',
            'longblob': 'byte[]',
        }

    def to_camel_case(self, snake_str):
        """将下划线命名转换为驼峰命名"""
        components = snake_str.split('_')
        return components[0] + ''.join(x.title() for x in components[1:])

    def get_java_type(self, mysql_type):
        """获取MySQL类型对应的Java类型"""
        # 提取基本类型（去除长度等信息）
        base_type = mysql_type.split('(')[0].lower()
        return self.type_mapping.get(base_type, 'String')

    def process_columns(self, columns):
        """处理列信息，转换数据类型和字段名"""
        processed_columns = []
        for column in columns:
            # 复制原始列信息
            processed_column = list(column)
            # 转换数据类型
            processed_column[1] = self.get_java_type(column[1])
            # 转换字段名为驼峰命名
            processed_column[0] = self.to_camel_case(column[0])

            # 为创建时间和修改时间字段添加特定注解
            if column[0] == 'created_time':
                processed_column.append(
                    '@JsonFormat(shape= JsonFormat.Shape.STRING,pattern="yyyy-MM-dd HH:mm:ss",timezone="GMT+8")\n\t@TableField(fill = FieldFill.INSERT)')
            elif column[0] == 'updated_time':
                processed_column.append(
                    '@JsonFormat(shape= JsonFormat.Shape.STRING,pattern="yyyy-MM-dd HH:mm:ss",timezone="GMT+8")\n\t@TableField(fill = FieldFill.UPDATE)')
            else:
                processed_column.append('')

            processed_columns.append(processed_column)
        return processed_columns

    def connect_db(self):
        """连接数据库"""
        try:
            self.connection = pymysql.connect(**self.db_config)
            return True
        except Exception as e:
            print(f"数据库连接失败: {str(e)}")
            return False

    def validate_table(self, table_name):
        """验证表是否存在"""
        if not self.connection:
            if not self.connect_db():
                return False

        try:
            with self.connection.cursor() as cursor:
                cursor.execute(f"SHOW TABLES LIKE '{table_name}'")
                return cursor.fetchone() is not None
        except Exception as e:
            print(f"验证表名失败: {str(e)}")
            return False

    def get_table_columns(self, table_name):
        """获取表的列信息"""
        if not self.connection:
            if not self.connect_db():
                return None

        try:
            with self.connection.cursor() as cursor:
                cursor.execute(f"SHOW FULL COLUMNS FROM {table_name}")
                columns = cursor.fetchall()
                return columns
        except Exception as e:
            print(f"获取表结构失败: {str(e)}")
            return None

    def generate_code(self, table_name, table_comment, generate_type):
        """生成代码"""
        if not self.validate_table(table_name):
            print(f"表 {table_name} 不存在")
            return False

        columns = self.get_table_columns(table_name)
        if not columns:
            return False

        # 处理表名，生成类名
        class_name = ''.join(word.capitalize() for word in table_name.split('_'))

        # 处理列信息
        processed_columns = self.process_columns(columns)

        # 根据生成类型选择模板
        if generate_type == 'java':
            self._generate_java_code(table_name, class_name, table_comment, processed_columns)
        elif generate_type == 'vue':
            self._generate_vue_code(table_name, class_name, table_comment, processed_columns)

        return True

    def _generate_java_code(self, table_name, class_name, table_comment, columns):
        """生成Java代码"""
        # 生成Model
        model_template = self.env.get_template('java/ClassifyModel.java')
        model_content = model_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment,
            columns=columns
        )

        # 生成Controller
        controller_template = self.env.get_template('java/ClassifyController.java')
        controller_content = controller_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment,
            columns=columns
        )

        # 生成Service接口
        service_template = self.env.get_template('java/ClassifyService.java')
        service_content = service_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment
        )

        # 生成Service实现类
        service_impl_template = self.env.get_template('java/ClassifyServiceImpl.java')
        service_impl_content = service_impl_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment
        )

        # 生成Mapper接口
        mapper_template = self.env.get_template('java/ClassifyMapper.java')
        mapper_content = mapper_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment
        )

        # 保存生成的文件
        self._save_file(f"{pro_root}/{root}/entity/model/{class_name}Model.java", model_content)
        self._save_file(f"{pro_root}/{root}/api/admin/{class_name}Controller.java", controller_content)
        self._save_file(f"{pro_root}/{root}/service/{class_name}Service.java", service_content)
        self._save_file(f"{pro_root}/{root}/service/serviceImpl/{class_name}ServiceImpl.java", service_impl_content)
        self._save_file(f"{pro_root}/{root}/mapper/{class_name}Mapper.java", mapper_content)

    def _generate_vue_code(self, table_name, class_name, table_comment, columns):
        """生成Vue代码"""
        # 生成API文件
        api_template = self.env.get_template('vue/api.js')
        api_content = api_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment
        )

        # 生成列表页面
        index_template = self.env.get_template('vue/index.page.vue')
        index_content = index_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment,
            columns=columns
        )

        # 生成表单页面
        form_template = self.env.get_template('vue/form.vue')
        form_content = form_template.render(
            class_name=class_name,
            table_name=table_name,
            table_comment=table_comment,
            columns=columns
        )

        # 保存生成的文件
        self._save_file(f"{vue_root}/src/api/{table_name}.js", api_content)
        self._save_file(f"{vue_root}/src/views/{table_name}/index.page.vue", index_content)
        self._save_file(f"{vue_root}/src/views/{table_name}/{table_name}-form.vue", form_content)

    def _save_file(self, file_path, content):
        """保存生成的文件"""
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)

    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()


def main():
    # 数据库配置
    db_config = {
        'host': 'localhost',
        'port': 3307,
        'user': 'root',
        'password': '123456',
        'database': 'drama',
        'charset': 'utf8mb4'
    }

    # 创建代码生成器实例
    generator = CodeGenerator(db_config)

    try:
        # 获取用户输入
        table_name = input("请输入表名: ")
        table_comment = input("请输入表注释: ")
        generate_type = input("请选择生成类型 (java/vue): ").lower()

        # 生成代码
        if generator.generate_code(table_name, table_comment, generate_type):
            print("代码生成成功！")
        else:
            print("代码生成失败！")

    finally:
        generator.close()


if __name__ == "__main__":
    # 项目根目录
    pro_root = "C:/Users/31100/Desktop/戏剧/drama-server"
    # 源码根目录 到src
    root = "/src/main/java/com/xjm"
    # vue的根目录
    vue_root = "C:/Users/31100/Desktop/戏剧/drama-admin"
    main()
